from app.db.models.tasks import Task
from app.handler.v2_jobs import JobHelper
from app.libs.enums import SECURITY_EVENT_PLUGIN_TUPLE
from app.db.models.asset_tasks import AssetTasks
from app.db.models.wass import ScanEvents
from app.db.models.jobs import Job
from app.libs.score import Score
from app.libs.enums import TaskType


class SumWarnings(object):
    level_score_map = {
        1: 1,
        2: 3,
        3: 5
    }
    level_map = {
        1: 'low',
        2: 'medium',
        3: 'high'
    }
    vul_category = ('低危', '中危', '高危')
    security_category = SECURITY_EVENT_PLUGIN_TUPLE
    ssl_category = ('certificate', 'protocol')
    content_category = ('keyword', 'privacy_disclosure')

    def __init__(self, warnings, task_type):
        self.warnings = warnings
        self.task_type = task_type

    def get(self):
        warnings = self.warnings
        task_type = self.task_type
        if task_type == 'vul':
            data = {'score': 0, 'low': 0, 'medium': 0, 'high': 0}
            return self.sum_warnings_vul(warnings, data)
        elif task_type == 'securityEvent':
            data = {'score': 0, **{k: 0 for k in self.security_category}}
            return self.sum_warnings_security(warnings, data)
        elif task_type == 'ssl':
            data = {'score': 0, 'certificate': 0, 'protocol': 0}
            return self.sum_warnings_ssl(warnings, data)
        elif task_type == 'content':
            data = {'score': 0, 'keyword': 0, 'privacy_disclosure': 0}
            return self.sum_warnings_content(warnings, data)
        return self.sum_warnings_all(warnings)

    def sum_warnings_all(self, warnings):
        vul = 50
        security_event = 50
        for warning in warnings:
            if self.is_warning_resolved(warning):
                continue
            _score = self.level_score_map.get(warning.get('level'), 0)
            if warning.get('levelName') in self.vul_category:
                vul -= _score
            elif warning.get('category') in self.ssl_category:
                vul -= _score
            elif warning.get('category') in self.security_category:
                security_event -= _score

        vul = vul if vul > 0 else 0
        security_event = security_event if security_event > 0 else 0
        return {'score': vul + security_event}

    def sum_warnings_vul(self, warnings, data):
        for warning in warnings:
            if self.is_warning_resolved(warning):
                continue
            category = self.level_map.get(warning.get('level'))
            if data.get(category) is not None:
                data[category] += 1
                data['score'] += 1

        return data

    def sum_warnings_security(self, warnings, data):
        for warning in warnings:
            if self.is_warning_resolved(warning):
                continue
            category = warning.get('category')
            if data.get(category) is not None:
                data[category] += 1
                data['score'] += 1

        return data

    def sum_warnings_content(self, warnings, data):
        for warning in warnings:
            if self.is_warning_resolved(warning):
                continue
            category = warning.get('category')
            if data.get(category) is not None:
                data[category] += 1
                data['score'] += 1
        return data

    def sum_warnings_ssl(self, warnings, data):
        return self.sum_warnings_security(warnings, data)

    @staticmethod
    def is_warning_resolved(warning):
        if warning.get('isResolved') is True or warning.get('ignoredTime'):
            return True
        return False


class LastTask(object):
    allow_key = ('endTime', 'taskType')

    def __init__(self, uid, task_type_list, request_args=None, job_query: dict = None, use_paginate: bool = False, only_score: bool = False):
        self.uid = uid
        self.task_type_list = task_type_list

        self.group = {}
        self.project = {}
        for key in self.allow_key:
            self.group[key] = {'$first': f'${key}'}
            self.project[key] = 1

        if not request_args:
            request_args = {}

        job_helper = JobHelper(request_args, uid, job_query, use_paginate, ('asset' in task_type_list))
        data = job_helper.get_data()
        self.job_count = job_helper.get_count()

        job_dict = {}
        task_ids = []
        for item in data:
            for task_type in task_type_list:
                new_task_type = f"{task_type}ResultId"
                if (task_id_list := getattr(item, new_task_type, None)) and task_id_list != [None]:
                    if not isinstance(task_id_list, list):
                        task_ids.append(task_id_list)
                    else:
                        task_ids.extend(task_id_list)

            job_dict[item.id] = item

        self.task_ids = task_ids
        self.job_dict = job_dict
        self.only_score = only_score
        self.data = None

    def get(self):
        data = {}
        for task_type in self.task_type_list:
            pipeline = [
                {'$match': {
                    'taskType': task_type, '_id': {'$in': self.task_ids}}},
                {'$sort': {'endTime': -1}},
                {'$group': {
                    '_id': '$jobId',
                    'warnings': {'$first': '$result.warnings'},
                    **self.group
                }},
                {'$project': {'_id': 1, 'warnings': {'$ifNull': ['$warnings', []]}, **self.project}}
            ]
            if self.only_score and task_type != TaskType.change_check.value:
                pipeline.append({'$project': {'_id': 1, 'warnings': {'category': 1, 'level': 1, }, **self.project}})
            data[task_type] = Task.objects.aggregate(pipeline, allowDiskUse=True)
        self.data = data
        return data

    def get_host_vul_warnings(self):
        return get_host_vul_warnings(self.uid)

    def get_warnings_group_task(self):
        job_id_task_dict = {}

        if not self.data:
            data = self.get()
        else:
            data = self.data

        for _ in data.values():
            for item in _:
                _job_id = item.get('_id')
                if _task := job_id_task_dict.get(_job_id):
                    _task['warnings'].extend(item.get('warnings', []))
                else:
                    if not (_job := self.job_dict.get(_job_id)):
                        continue
                    item['note'] = getattr(_job, 'note', '')
                    item['job_id'] = str(_job_id)
                    item['target'] = getattr(_job, 'targetUrl', '')
                    item['sourceIp'] = getattr(_job, 'sourceIp', '')
                    job_id_task_dict[_job_id] = item

        return tuple(job_id_task_dict.values())


def get_host_vul_warnings(uid, job_ids=None):
    query_dict = {"uid": uid, "verification.isVerified": True, "assetType": "host"}
    if job_ids:
        query_dict["_id"] = {"$in": job_ids}
    jobs = Job.objects.find(query_dict)
    job_warning_maps = dict()
    total = 0
    job_id_list = []
    for job in jobs:
        job_id_list.append(job.pk)
        job_warning_maps[job.pk] = {"targetUrl": job.targetUrl, 'note': job.note, "warnings": [], "has_task": False,
                                    "task": dict()}

    severityMap = Score.severityMap
    aggregate = [
        {'$match': {'uid': uid, 'taskType': "hostVul", "jobId": {"$in": job_id_list}, "hostVulResultId": {"$exists": True}}},
        {'$group': {'_id': '$jobId', 'id': {'$last': '$_id'}, 'hostVulResultId': {'$last': '$hostVulResultId'}}},
        {'$lookup': {'from': "tasks", 'localField': "hostVulResultId", 'foreignField': "_id", 'as': "task"}},
    ]
    asset_tasks = AssetTasks.objects.aggregate(aggregate)
    for asset_task in asset_tasks:
        if not asset_task.get('task'):
            continue

        task = asset_task.get('task')[0]
        job_id = task.get('jobId')
        job_warning_maps[job_id]['has_task'] = True
        job_warning_maps[job_id]['task']["asset_task_id"] = str(asset_task.get("id"))
        job_warning_maps[job_id]['task']["task_id"] = str(task.get("_id"))
        job_warning_maps[job_id]['task']["trigger_type"] = str(task.get("triggerType"))

        query = {"task_session_id": task.get("taskSessionId"), "event_name": "network_vuln"}
        host_vul_events = ScanEvents.objects.find(query)
        for event in host_vul_events:
            item = event.detail
            warning = {"level": severityMap.get(item["severity"], {}).get("level")}
            total += 1
            job_warning_maps[task.get('jobId')]['warnings'].append(warning)

    return total, job_warning_maps
